#!/usr/bin/env python
# encoding: utf-8
"""
manageovz.py

Created by Steven McGrath on 2009-06-01.
Copyright (c) 2009 __MyCompanyName__. All rights reserved.
"""

import sys
import os
import cmd
import openvz
import readline
import commands

class CLI(cmd.Cmd):
  def __init__(self):
    cmd.Cmd.__init__(self)
    self.prompt = 'ovzmanager> '
  
  # Command Help Block
  def help_help(self):
    print '\n'.join(['syntax: help [topic]',
                     '  Display Help information.'
                    ])
  def help_create(self, arg):
    print '\n'.join(['syntax: create [node]',
                     '  Create a new VPS'
                    ])
  def help_destroy(self):
    print '\n'.join(['syntax: destroy [node]',
                     '  Remove the specified node from the environment.'
                     '\n  NOTE:  This cannod be undone.'
                    ])
  def help_start(self, arg):
    print '\n'.join(['syntax: start [node]',
                     '  Start up the specified VPS'
                    ])
  def help_stop(self):
    print '\n'.join(['syntax: stop [node]',
                     '  Stop the specified VPS.'
                    ])
  def help_restart(self, arg):
    print '\n'.join(['syntax: restart [node]',
                     '  Restart the specified VPS'
                    ])
  def help_chmem(self):
    print '\n'.join(['syntax: chmem [node]',
                     '  Change the memory allocations for the specified VPS.'
                    ])
  def help_chdisk(self, arg):
    print '\n'.join(['syntax: chdisk [node]',
                     '  Change the disk allocations for the specified VPS'
                    ])
  def help_status(self):
    print '\n'.join(['syntax: status [node]',
                     '  Display status information for all of the nodes, or just the specified node.'
                    ])
  def help_shell(self, arg):
    print '\n'.join(['syntax: shell [node]',
                     '  initiates a shell on the specified node.'
                    ])
  def help_quit(self, arg):
    print '\n'.join(['syntax: quit',
                     '  Quits OVZManager'
                    ]) 
  
  # Command Completions Block
  def complete_destroy(self, text, line, begidx, endidx):
    return complete()
  def complete_stop(self, text, line, begidx, endidx):
    return complete()
  def complete_start(self, text, line, begidx, endidx):
    return complete()
  def complete_restart(self, text, line, begidx, endidx):
    return complete()
  def complete_status(self, text, line, begidx, endidx):
    return complete()
  def complete_chmem(self, text, line, begidx, endidx):
    return complete()    
  def complete_chdisk(self, text, line, begidx, endidx):
    return complete()
  def complete_shell(self, text, line, begidx, endidx):
    return complete()


  # new command
  def do_create(self, arg):
    
    dns_server = '172.16.0.1'
    dns_domain = 'cugnet.local'
    
    build_vps = False
    node = {
             'hostname': '',
                 'veid': '',
                 'user': 'none',
               'memory': 50,
                 'disk': 1,
             'template': 'ubuntu-9.04-x86_64',
            'ipaddress': '',
               'status': 'not created'
           }
    
    while build_vps == False:
      name    = raw_input("    [%20s] Name : " % node['hostname'])
      veid    = raw_input("                   [%5s] VEID : " % node['veid'])
      memory  = raw_input("              [%5s MB] Memory : " % node['memory'])
      disk    = raw_input("                [%5s GB] Disk : " % node['disk'])
      ipadd   = raw_input("      [%12s] IP Address : " % node['ipaddress'])
      
      print " Available Templates\n--------------------"
      for item in os.listdir('/srv/vz/template/cache'):
        print item.split('.')[0]
      
      temp    = raw_input("[%20s] Template : " % node['template'])
      
      if name == 'break' or veid == 'break' or memory == 'break' or disk == 'break' or ipadd == 'break' or temp == 'break':
        return
      
      if name != '':
        node['hostname']  = name
      if veid != '':
        node['veid']      = veid
      if memory != '':
        node['memory']    = memory
      if disk != '':
        node['disk']      = disk
      if temp != '':
        node['template']  = temp
      if ipadd != '':
        node['ipaddress'] = ipadd
      
      if node['hostname'] == '':
        print "\nERROR: Need a Valid Name!\n"
        continue
      if node['veid'] == '':
        print "\nERROR: Need a VEID Number!\n"
        continue
      if node['ipaddress'] == '':
        print "\nERROR: Need a IP Address!\n"
        continue
      if float(node['memory']) < 1:
        print "\nERROR: No memory Allocated!\n"
        continue
      if float(node['disk']) < 1:
        print "\nERROR: Insufficient Disk Allocated!\n"
        continue
      
      openvz.print_node(node)
      confirm = raw_input('Are these settings ok? : ')
      
      if confirm == 'yes':
        break
    
    inodes = openvz.inodes(node['disk'])
    disk_soft = openvz.gb2kb(node['disk'])
    disk_hard = openvz.gb2kb(float(node['disk']) * 1.10)
    guar_mem_soft = openvz.mb2pages(node['memory'])
    guar_mem_hard = openvz.mb2pages(float(node['memory']) * 1.10)
    burst_mem_soft = openvz.mb2pages(float(node['memory']) * 1.50)
    burst_mem_hard = openvz.mb2pages(float(node['memory']) * 1.60)
    kern_mem_soft = openvz.mb2bytes(float(node['memory']) * 0.15)
    kern_mem_hard = openvz.mb2bytes(float(node['memory']) * 0.20)
    
    config = ' '.join(["--hostname %s" % node['hostname'],
                       "--name %s" % node['hostname'],
                       "--nameserver %s" % dns_server,
                       "--searchdomain %s" % dns_domain,
                       "--ipadd %s" % node['ipaddress'],
                       "--diskspace %s:%s" % (disk_soft, disk_hard),
                       "--diskinodes %s" % inodes,
                       "--privvmpages %s:%s" % (burst_mem_soft, burst_mem_hard),
                       "--oomguarpages %s:%s" % (guar_mem_soft, guar_mem_hard),
                       "--kmemsize %s:%s" % (kern_mem_soft, kern_mem_hard),
                       "--save"
                      ])
    os.system("vzctl create %s --ostemplate %s" % (node['veid'], node['template']))
    os.system("vzctl set %s %s" % (node['veid'], config))
    os.system("vzctl start %s" % node['veid'])
        
  
  # destroy command
  def do_destroy(self, arg):
    node = openvz.node(arg)
    if node['veid'] == 'none':
      return
    else:
      openvz.print_node(node)
      print '\n'.join(['!!!! WARNING !!!! : This will remove the node from the system.',
                       '                    This action is unrecoverable!'
                      ])
      confirm = raw_input('\nAre you sure you want to destroy this node? : ')
      if confirm.upper() != "YES":
        return
        
      if node['status'] == 'running':
        os.system("vzctl stop %s" % node['veid'])
      os.system("vzctl destroy %s" % node['veid'])
  
  # shell command
  def do_shell(self, arg):
    node = openvz.node(arg)
    if node['veid'] == 'none':
      return
    else:
      os.system('vzctl enter %s' % node['veid'])
  
  def do_migrate(self, arg):
    node = openvz.node(arg.split()[0])
    user = arg.split()[1]
    
    if node['veid'] == 'none':
      return
    else:
      print "Are you sure you want to migrate %s to the following node : "
      openvz.print_node(node)
      yesno = raw_input('(yes/no) :> ')
      
      if yesno.lower() == 'yes':
        # check to make sure user exists and pull user's home
        # modify user in ldap to use /home
        # move data to VPS
        # remove data from the shell server.
        pass
  
  def do_newuser(self, arg):
    # ask for user information
    # create user in ldap
    pass
    
  def do_backup(self, arg):
    # Backup the VPS
    pass
  
  # start command
  def do_start(self, arg):
    node = openvz.node(arg)
    if node['veid'] == 'none':
      return
    else:
      print "Starting the Following Node..."
      openvz.print_node(node)
      os.system("vzctl start %s" % node['veid'])
    
  # stop command
  def do_stop(self, arg):
    node = openvz.node(arg)
    if node['veid'] == 'none':
      return
    else:
      print "Stopping the Following Node..."
      openvz.print_node(node)
      os.system("vzctl stop %s" % node['veid'])

  # restart command
  def do_restart(self, arg):
    node = openvz.node(arg)
    if node['veid'] == 'none':
      return
    else:
      print "Restarting the Following Node..."
      openvz.print_node(node)
      os.system("vzctl restart %s" % node['veid'])

  # chmem command
  def do_chmem(self, arg):
    node = openvz.node(arg)    
    if node['veid'] == 'none':
      return
    else:
      openvz.print_node(node)
      new_mem = raw_input('\nEnter new Memory size in Megabytes : ')
      if new_mem == "":
        return
      guar_mem_soft = openvz.mb2pages(new_mem)
      guar_mem_hard = openvz.mb2pages(float(new_mem) * 1.10)
      burst_mem_soft = openvz.mb2pages(float(new_mem) * 1.50)
      burst_mem_hard = openvz.mb2pages(float(new_mem) * 1.60)
      kern_mem_soft = openvz.mb2bytes(float(new_mem) * 0.15)
      kern_mem_hard = openvz.mb2bytes(float(new_mem) * 0.20)
      
      config = ' '.join(["--privvmpages %s:%s" % (burst_mem_soft, burst_mem_hard),
                         "--oomguarpages %s:%s" % (guar_mem_soft, guar_mem_hard),
                         "--kmemsize %s:%s" % (kern_mem_soft, kern_mem_hard),
                         "--save"
                        ])
      os.system("vzctl set %s %s" % (node['veid'], config))

  # chdisk command  
  def do_chdisk(self, arg):
    node = openvz.node(arg)    
    if node['veid'] == 'none':
      return
    else:
      openvz.print_node(node)
      new_disk = raw_input('\nEnter new Disk size in Gigabytes : ')
      if new_disk == "":
        return
      inodes = openvz.inodes(new_disk)
      disk_soft = openvz.gb2kb(new_disk)
      disk_hard = openvz.gb2kb(float(new_disk) * 1.10)
      
      config = ' '.join(["--diskspace %s:%s" % (disk_soft, disk_hard),
                         "--diskinodes %s" % inodes,
                         "--save"
                        ])
      os.system("vzctl set %s %s" % (node['veid'], config))

  # status command
  def do_status(self, arg):
    if arg == "":
      nodes = openvz.nodes()
      print " VEID    Node Name       User     RAM   Disk   Status"
      print "----- --------------- ---------- ----- ----- ----------"
      for node in nodes:
        if node['status'] == 'running':
          print "%5s %15s %10s %5s %5s \033[32m%10s\033[0m" % (node['veid'], node['hostname'], node['user'], node['memory'], node['disk'], node['status'])
        if node['status'] == 'stopped':
          print "%5s %15s %10s %5s %5s \033[31m%10s\033[0m" % (node['veid'], node['hostname'], node['user'], node['memory'], node['disk'], node['status'])
    else:
      openvz.print_node(openvz.node(arg))
  
  # quit command
  def do_quit(self, arg):
    sys.exit(1)

def complete():
  ret_val = []
  for node in openvz.nodes():
    ret_val.append(node['hostname'])
    ret_val.append(node['veid'])
  return re_val
  
    
def main():
  cli = CLI()
  cli.cmdloop()

if __name__ == '__main__':
  main()

